import os.path

from infrastructure.constant import file, models
from infrastructure.utils import prompt, text
from application.base.ollama_server import OllamaServer


class EntityGoCodeGenerator:

    def __init__(self):
        self.server = OllamaServer()

    def generate(self, query: str = '快速排序算法代码示例', model_name: str = models.DEFAULT_MODEL) -> str:
        result = self.server.invoke(
            model_name=model_name,
            messages=prompt.append_user_message(query)
        )
        # 去除深度思考tag
        if "<think>" in result:
            result = text.remove_think_blocks(result).strip()
        # 去除go-markdown-code包裹
        if "```go" in result:
            result = text.extract_md_go_code_blocks(result).strip()
        file_path = os.path.join(file.GENERATE_GO_CODE_PATH, file.FILE_NAME_OF_ENTITY)
        with open(file_path, mode='w', encoding='utf-8') as f:
            f.write(result)
        return result
